#include "limbo/errc.h"
#include "limbo/ip/state.hpp"
#include "limbo/stack.hpp"
#include "limbo/udp/state.hpp"
#include "test-utils.h"

using namespace limbo;

using IP = ip::State<void, ip::Opts::calculate_checksum>;
using UDP = udp::State<IP>;
using MyStack = Stack<Layer<UDP>, Layer<IP>>;

unsigned char example_raw[20 + 8 + 5] = {
    /* hdr */ 0x45, 0x00, 0x00, 0x21, 0x00, 0x05, 0x00, 0x00,
    /* ttl */ 0x06, 0x11, 0xaa, 0x50, 0x0a, 0x00, 0x1f, 0x7c,
    /* dst */ 0xe0, 0x00, 0x00, 0xfb, 0x00, 0x07, 0x00, 0x08,
    /* len */ 0x00, 0x0D, 0x00, 0x00, 'h',  'e',  'l',  'l',  'o'};
auto example_chunk = Chunk(example_raw, sizeof(example_raw));
auto payload_chunk = Chunk((void *)(example_raw + 28), 5);

auto src = make_address("10.0.31.124");
auto dst = make_address("224.0.0.251");

TEST_CASE("send", "[ip+udp]") {
  auto stack = MyStack();
  auto &ip_state = stack.get<1, 0>();
  auto &udp_state = stack.get<0, 0>();
  ip_state.init(0x04);

  auto ip_ctx = IP::Context{src, dst, 6, ip::Proto::udp, 0, 0, nullptr};
  auto udp_ctx = UDP::Context{7, 8, &ip_ctx};
  unsigned char buff[sizeof(example_raw)];

  SECTION("no enough space") {
    auto zero_chunk = Chunk(buff, 0);
    auto result = stack.send(zero_chunk, udp_ctx, payload_chunk);
    REQUIRE(!result);
    CHECK(result.state() == &udp_state);
    CHECK(result.error_code() == (uint32_t)Errc::no_enough_space);
  }

  auto buff_chunk = Chunk(buff, sizeof(buff));

  SECTION("send with payload copying") {
    auto result = stack.send(buff_chunk, udp_ctx, payload_chunk);

    REQUIRE(result);
    auto res_buff = result.consumed();
    CHECK((void *)res_buff.data() == (void *)buff);
    CHECK(res_buff.size() == example_chunk.size());
    CHECK(res_buff == example_chunk);
  }

  SECTION("send envelope") {
    auto sender = stack.send_envelope(buff_chunk, udp_ctx, payload_chunk);
    REQUIRE(sender);
    CHECK(sender.header().size() == 20 + 8);
    auto header = Chunk(example_raw, sender.header().size());
    CHECK(sender.header() == header);
    CHECK(sender.trailer() == Chunk());
  }
}

TEST_CASE("receive", "[ip+udp]") {
  auto stack = MyStack();
  auto &udp_state = stack.get<0, 0>();
  auto result = stack.recv(example_chunk, nullptr);
  REQUIRE(result);
  REQUIRE(result.consumed() == example_chunk);
  REQUIRE(result.state() == &udp_state);
  auto &p = udp_state.get_parsed();
  CHECK(p.source_port == 7);
  CHECK(p.dest_port == 8);
  CHECK(p.length == payload_chunk.size() + 8);
  CHECK(p.payload == payload_chunk);
  CHECK(p.container->source == src);
  CHECK(p.container->destination == dst);
}

TEST_CASE("icmp packet is ignored", "[ip+udp]") {
  unsigned char data[] = {
      0x45, 0xc0, 0x00, 0x48, 0x95, 0x78, 0x00, 0x00, 0x40, 0x01, 0x91, 0x67,
      0x0a, 0x00, 0x1f, 0x92, 0x0a, 0x00, 0x1f, 0x84, 0x03, 0x03, 0x50, 0x3c,
      0x00, 0x00, 0x00, 0x00, 0x45, 0x00, 0x00, 0x2c, 0x8e, 0x6a, 0x00, 0x00,
      0x80, 0x11, 0x59, 0x41, 0x0a, 0x00, 0x1f, 0x84, 0x0a, 0x00, 0x1f, 0x92,
      0xcc, 0xc3, 0x21, 0xa2, 0x00, 0x18, 0x17, 0x64, 0x4d, 0x46, 0x4e, 0x50,
      0x02, 0x01, 0x00, 0x00, 0x09, 0x47, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  };
  auto data_chunk = Chunk(data, sizeof(data));
  auto stack = MyStack();
  auto &ip_state = stack.get<1, 0>();
  auto result = stack.recv(data_chunk, nullptr);
  REQUIRE(result);
  CHECK(result.state() == &ip_state);
}

TEST_CASE("checksum validation", "[ip+udp]") {
  using IP = ip::State<void, ip::Opts::calculate_checksum |
                                 ip::Opts::validate_checksum>;
  using UDP = udp::State<IP, udp::Opts::calculate_checksum |
                                 udp::Opts::validate_checksum>;
  using MyStack = Stack<Layer<UDP>, Layer<IP>>;

  unsigned char example_raw[] = {
      0x45, 0x00, 0x00, 0x26, 0x46, 0xf4, 0x40, 0x00, 0x40, 0x11,
      0xaa, 0x6d, 0xc0, 0xa8, 0x64, 0x0f, 0xc0, 0xa8, 0x64, 0x05,
      0xca, 0x19, 0xfc, 0xe2, 0x00, 0x12, 0x0f, 0x66, 0x54, 0xf7,
      0x86, 0x72, 0x78, 0x32, 0xbf, 0x5e, 0xcd, 0x06,
  };
  auto example_chunk = Chunk(example_raw, sizeof(example_raw));
  auto stack = MyStack();
  auto &ip_state = stack.get<1, 0>();

  auto &udp_state = stack.get<0, 0>();
  auto result = stack.recv(example_chunk, nullptr);
  REQUIRE(result);
  REQUIRE(result.consumed() == example_chunk);
  REQUIRE(result.state() == &udp_state);

  auto src = make_address("192.168.100.15");
  auto dst = make_address("192.168.100.5");
  ip_state.init(0x46f3);
  auto ip_ctx = IP::Context{src, dst, 64, ip::Proto::udp, 0b010, 0, nullptr};
  auto udp_ctx = UDP::Context{51737, 64738, &ip_ctx};
  unsigned char buff_raw[sizeof(example_raw)];
  auto buff_chunk = Chunk(buff_raw, sizeof(buff_raw));

  auto payload_sz = 10;
  auto payload_off = example_chunk.size() - payload_sz;
  auto payload_chunk = Chunk(example_raw + payload_off, payload_sz);
  result = stack.send(buff_chunk, udp_ctx, payload_chunk);
  REQUIRE(result);
  CHECK(buff_chunk == example_chunk);
}
